# file:train_test.py
# 将先前编写好的计算准确率方法导入:
from CNN import cal_correction
def train_model(model, train_data, val_data, batch, epoch, loss_func, opt):
	""" @参数说明:
		model - 需要训练的模型;
		train_data - 训练数据;
		val_data - 验证数据;
		batch - 每次训练的批量;
		epoch - 训练轮数;
		loss_func - 损失函数;
		opt - 优化器."""
	for i in range(epoch):
		# 定义记录损失值、精度的容器:
		train_corrections = []
		train_losses = []
		for idx, (img, label) in enumerate(train_data):
			# 数据处理:
			# 转化成cuda类型：
			img, label = img.to('cpu'), label.to('cpu')
			# 保留梯度:
			img = img.clone().requires_grad_(True)
			label = label.clone().detach()
			""" ----- 前向传播 ----- """ 
			# 设置模型为训练模式:
			model.train()
			# 将数据喂入网络:
			output = model(img)
			# 计算精度:
			train_acc = cal_correction(output, label)
			train_corrections.append(train_acc)
			# 计算损失值:
			train_loss = loss_func(output, label)
			train_losses.append(train_loss)
			# 清空优化器梯度:
			opt.zero_grad()
			""" ----- 反向传播 ----- """
			# 反传损失值:
			train_loss.backward()
			# 用优化器梯度下降:
			opt.step()
			
			if idx % (batch * 100) == 0:
				"""每训练好100批次验证并展示训练结果"""
				model.eval()
				val_record = []
				for (data, target) in val_data:
					# 转化数据类型:
					data, target = data.to('cpu'), target.to('cpu')
					data, target = data.clone().requires_grad_(True), target.clone().detach()
					# 将数据喂入:
					out = model(data)
					# 预测准确率:
					val_acc = cal_correction(out, target)
					val_record.append(val_acc)
					# 打印训练、验证结果:
					print(f'epoch{i+1}: Train Acc={train_corrections[-1]} Train Loss={train_losses[-1]} Val Acc={val_record[-1]}')			

def test_model(model, test_data):
	"""测试模型"""
	# 定义记录测试精度的容器:
	test_acc = []
	for idx, (img, label) in enumerate(test_data):
		# 处理数据:
		img, label = img.to('cpu'), label.to('cpu')
		img, label = img.clone().requires_grad_(True), label.clone().detach()
		# 投入数据:
		output = model(img)
		# 比对结果:
		acc = cal_correction(output, label)
		test_acc.append(acc)
		print(f'Test Acc={test_acc[-1]}%')
